#include <ntifs.h>
#include <ntstrsafe.h>

#define DRIVERNAME "FSFltr"
#define PAGEDCODE code_seg("PAGE")
#define LOCKEDCODE code_seg()
#define INITCODE code_seg("INIT")

#define RECORD_LOG_IRP 0
#define RECORD_LOG_FASTIO 1

#define IS_MY_DEVICE_OBJECT(_devObj) \
    (((_devObj) != NULL) && \
	((_devObj)->DriverObject == GlobalData->DriverObject) && \
      ((_devObj)->DeviceExtension != NULL))

#define IS_MY_CONTROL_DEVICE_OBJECT(_devObj) \
	(((_devObj) != NULL && (_devObj)->DriverObject == GlobalData->DriverObject) &&\
	((_devObj) == GlobalData->ControlDeviceObject))

# ifdef _X86_
  # define MAX_FAST_REFS 7
# else
  # define MAX_FAST_REFS 15
# endif

typedef enum _FASTIO_TYPE {
	READ, WRITE, CHECK_IF_POSSIBLE, QUERY_BASIC_INFO, QUERY_STANDART_INFO, IO_LOCK, UNLOCK_SINGLE, UNLOCK_ALL, UNLOCK_ALL_BY_KEY, DEVICE_CONTROL, DETACH_DEVICE,
	NETWORK_OPEN_INFO, MDL_READ, MDL_READ_COMPLETE, PREPARE_MDL_WRITE, MDL_WRITE_COMPLETE, READ_COMPRESSED, WRITE_COMPRESSED, MDL_READ_COMPLETE_COMPRESSED,
	MDL_WRITE_COMPLETE_COMPRESSED, QUERY_OPEN } FASTIO_TYPE, *PFASTIO_TYPE;

typedef struct _DEVICE_EXTENSION {
	PDEVICE_OBJECT LowDevice;
	WCHAR VolumeName[20];
	KEVENT Event;
} DEVICE_EXTENSION, *PDEVICE_EXTENSION;

typedef enum _FILTER_VALUE_TYPE {
	FILTER_PROCESS_TYPE, FILTER_FILE_TYPE, FILTER_FASTIO_TYPE, 
	FILTER_IRP_MJ_TYPE, FILTER_IRP_STATUS 
} FILTER_VALUE_TYPE, *PFILTER_VALUE_TYPE;

typedef struct _FILTER_VALUE 
{
	ULONG Size;
	BOOLEAN UseIrpMj;
	UCHAR FilterIrpCount;
	ULONG *FilterIrpMj;
	BOOLEAN UseFastIo;
	UCHAR FilterFastIoCount;
	ULONG *FilterFastIo;
	BOOLEAN UseStatus;
	ULONG Status;
	BOOLEAN UseProcName;
	BOOLEAN UseFilePath;
	PVOID Next;
	UNICODE_STRING FilePath;
	UNICODE_STRING ProcName;
} FILTER_VALUE, *PFILTER_VALUE;

typedef struct _CONTROL_DEVICE_EXTENSION {
	ULONG CountConnections;
	PDRIVER_OBJECT DriverObject;
} CONTROL_DEVICE_EXTENSION, *PCONTROL_DEVICE_EXTENSION;

typedef struct _RECORD_IRP {
	LARGE_INTEGER CompletionTime;

	UCHAR IrpMj;
	UCHAR IrpMn;
	ULONG IrpFlags;
	NTSTATUS ReturnStatus;
	ULONG ThreadID;
} RECORD_IRP, *PRECORD_IRP;

typedef struct _RECORD_FASTIO
{
	LARGE_INTEGER CompletionTime;
	ULONG ThreadId;
	FASTIO_TYPE Action;
	NTSTATUS ReturnStatus;
} RECORD_FASTIO, *PRECORD_FASTIO;

typedef union _RECORD_IO {
	RECORD_IRP RecordIrp;
	RECORD_FASTIO RecordFastIo;
} RECORD_IO, *PRECORD_IO;

typedef struct _LOG_RECORD
{
	WCHAR FilePath[260];
	WCHAR ProcessName[260];
	RECORD_IO Record;
	UCHAR Type;
	PVOID Next;
} LOG_RECORD, *PLOG_RECORD;

typedef struct _GLOBAL_DATA {
	PDRIVER_OBJECT DriverObject;
	PDEVICE_OBJECT ControlDeviceObject;
	FAST_MUTEX Mutex;
	FAST_MUTEX MutexProc;
	KSPIN_LOCK lLock;
	KEVENT MountedEvent;
	PLOG_RECORD LogRecord;
	PLOG_RECORD CurrentLogRecord;
	BOOLEAN Monitoring;
	ULONG LogRecordCount;
	PFILTER_VALUE FilterValue;
	PFILTER_VALUE CurrentFilterValue;
	BOOLEAN Filtering;
} GLOBAL_DATA, *PGLOBAL_DATA;

PGLOBAL_DATA GlobalData;

typedef struct _CONTROL_AREA
{
     PVOID Segment;
     LIST_ENTRY DereferenceList;
     ULONG Counters[5];
     ULONG u;
     PFILE_OBJECT FilePointer;
} CONTROL_AREA, *PCONTROL_AREA;

typedef struct _CONTROL_AREA_VISTA
{
    PVOID Segment;
    LIST_ENTRY DereferenceList;
    ULONG_PTR Counters [4];
    ULONG Flags [2];
    PFILE_OBJECT FilePointer;
} CONTROL_AREA_VISTA, *PCONTROL_AREA_VISTA;

typedef struct _SEGMENT {
    PCONTROL_AREA                ControlArea; // 0x00
} SEGMENT, * PSEGMENT;

typedef struct _SEGMENT_VISTA {
	PCONTROL_AREA_VISTA ControlArea;
} SEGMENT_VISTA, *PSEGMENT_VISTA;

typedef struct _MMADDRESS_NODE {
    union
    {
        LONG_PTR Balance : 2;
        PVOID Parent;
    }
    u1;

    PVOID LeftChild;
    PVOID RightChild;
    ULONG_PTR StartingVpn;
    ULONG_PTR EndingVpn;
} MMADDRESS_NODE,
* PMMADDRESS_NODE;

typedef struct _SECTION_OBJECT {
    MMADDRESS_NODE                Address; // 0x00
    PSEGMENT                    Segment; // 0x14 PTR _SEGMENT
    LARGE_INTEGER                SizeOfSection; // 0x18
    ULONG                        Flags; // 0x20 MMSECTION_FLAGS
    ULONG                        InitialPageProtection; // 0x24
} SECTION_OBJECT, *PSECTION_OBJECT;

typedef struct _SECTION_OBJECT_VISTA {
    MMADDRESS_NODE                Address; // 0x00
    PSEGMENT_VISTA	               Segment; // 0x14 PTR _SEGMENT
    LARGE_INTEGER                SizeOfSection; // 0x18
    ULONG                        Flags; // 0x20 MMSECTION_FLAGS
    ULONG                        InitialPageProtection; // 0x24
} SECTION_OBJECT_VISTA, *PSECTION_OBJECT_VISTA;

typedef struct _OBJECT_DIRECTORY_INFORMATION
{
    UNICODE_STRING ObjectName;
    UNICODE_STRING ObjectTypeName;
}OBJECT_DIRECTORY_INFORMATION, *POBJECT_DIRECTORY_INFORMATION;

extern POBJECT_TYPE *MmSectionObjectType;

NTSTATUS AttachDeviceToDeviceStack(PDEVICE_OBJECT TargetDeviceObject, PDEVICE_OBJECT Vdo);
NTSTATUS DetachDeviceFromDeviceStack(PDEVICE_OBJECT TargetDeviceObject);

NTSTATUS DispatchAny(PDEVICE_OBJECT FiDO, PIRP Irp);
NTSTATUS DispatchMonitoring(PDEVICE_OBJECT FiDO, PIRP Irp);
NTSTATUS DispatchFsControl(PDEVICE_OBJECT FiDO, PIRP Irp);
VOID DriverUnload(PDRIVER_OBJECT DriverObject);
NTSTATUS CdoDispatchCreate(PDEVICE_OBJECT FiDO, PIRP Irp);
NTSTATUS CdoDispatchClose(PDEVICE_OBJECT FiDO, PIRP Irp);
NTSTATUS CdoDispatchDeviceControl(PDEVICE_OBJECT FiDO, PIRP Irp, PULONG Info);
BOOLEAN IsAttachedToDevice(PDEVICE_OBJECT AttachedDevice);

NTSTATUS AnyCompletionRoutine(PDEVICE_OBJECT FiDO, PIRP Irp, PVOID Context);
NTSTATUS MountedCompletionRoutine(PDEVICE_OBJECT FiDO, PIRP Irp, PVOID Context);
NTSTATUS FileNameQueryComplete(PDEVICE_OBJECT DeviceObject, PIRP Irp, PVOID Context);
NTSTATUS CompleteRequest(PIRP Irp, NTSTATUS Status, ULONG_PTR Info);

BOOLEAN FastIoCheckIfPossible(PFILE_OBJECT FileObject, PLARGE_INTEGER FileOffset, ULONG Length, BOOLEAN Wait,
						 ULONG LockKey, BOOLEAN CheckForReadOperation, PIO_STATUS_BLOCK IoStatus, PDEVICE_OBJECT DeviceObject);
BOOLEAN FastIoRead(PFILE_OBJECT FileObject, PLARGE_INTEGER FileOffset, ULONG Length, BOOLEAN Wait, 
				   ULONG LockKey, PVOID Buffer, PIO_STATUS_BLOCK IoStatus, PDEVICE_OBJECT DeviceObject);
BOOLEAN FastIoWrite(PFILE_OBJECT FileObject, PLARGE_INTEGER FileOffset, ULONG Length, BOOLEAN Wait,
			   ULONG LockKey, PVOID Buffer, PIO_STATUS_BLOCK IoStatus, PDEVICE_OBJECT DeviceObject);
BOOLEAN FastIoQueryBasicInfo(PFILE_OBJECT FileObject, BOOLEAN Wait, PFILE_BASIC_INFORMATION Buffer, 
							 PIO_STATUS_BLOCK IoStatus, PDEVICE_OBJECT DeviceObject);
BOOLEAN FastIoQueryStandardInfo(PFILE_OBJECT FileObject, BOOLEAN Wait, PFILE_STANDARD_INFORMATION Buffer, 
								PIO_STATUS_BLOCK IoStatus, PDEVICE_OBJECT DeviceObject);
BOOLEAN FastIoLock(PFILE_OBJECT FileObject, PLARGE_INTEGER FileOffset, PLARGE_INTEGER Length, PEPROCESS ProcessId,
			  ULONG Key, BOOLEAN FailImmediately, BOOLEAN ExclusiveLock, PIO_STATUS_BLOCK IoStatus, PDEVICE_OBJECT DeviceObject);
BOOLEAN FastIoUnlockSingle(PFILE_OBJECT FileObject, PLARGE_INTEGER FileOffset, PLARGE_INTEGER Length,
					  PEPROCESS ProcessId, ULONG Key, PIO_STATUS_BLOCK IoStatus, PDEVICE_OBJECT DeviceObject);
BOOLEAN FastIoUnlockAll(PFILE_OBJECT FileObject, PEPROCESS ProcessId, PIO_STATUS_BLOCK IoStatus, PDEVICE_OBJECT DeviceObject);
BOOLEAN FastIoUnlockAllByKey(PFILE_OBJECT FileObject, PVOID ProcessId, ULONG Key, PIO_STATUS_BLOCK IoStatus, PDEVICE_OBJECT DeviceObject);
BOOLEAN FastIoDeviceControl(PFILE_OBJECT FileObject, BOOLEAN Wait, PVOID InputBuffer, ULONG InputBufferLength, PVOID OutputBuffer,
					   ULONG OutputBufferLength, ULONG IoControlCode, PIO_STATUS_BLOCK IoStatus, PDEVICE_OBJECT DeviceObject);
VOID FastIoDetachDevice (PDEVICE_OBJECT SourceDevice,PDEVICE_OBJECT TargetDevice);
BOOLEAN FastIoQueryNetworkOpenInfo (PFILE_OBJECT FileObject, BOOLEAN Wait, PFILE_NETWORK_OPEN_INFORMATION Buffer, 
							   PIO_STATUS_BLOCK IoStatus, PDEVICE_OBJECT DeviceObject);

BOOLEAN FastIoMdlRead(PFILE_OBJECT FileObject, PLARGE_INTEGER FileOffset, ULONG Length,
					  ULONG LockKey, PMDL *MdlChain, PIO_STATUS_BLOCK IoStatus, PDEVICE_OBJECT DeviceObject);
BOOLEAN FastIoMdlReadComplete(PFILE_OBJECT FileObject, PMDL MdlChain, PDEVICE_OBJECT DeviceObject);
BOOLEAN FastIoPrepareMdlWrite(PFILE_OBJECT FileObject, PLARGE_INTEGER FileOffset, ULONG Length,
							   ULONG LockKey, PMDL *MdlChain, PIO_STATUS_BLOCK IoStatus, PDEVICE_OBJECT DeviceObject);
BOOLEAN FastIoMdlWriteComplete(PFILE_OBJECT FileObject, PLARGE_INTEGER FileOffset, PMDL MdlChain, PDEVICE_OBJECT DeviceObject);
BOOLEAN FastIoReadCompressed(PFILE_OBJECT FileObject, PLARGE_INTEGER FileOffset, ULONG Length,
							 ULONG LockKey, PVOID Buffer, PMDL *MdlChain, PIO_STATUS_BLOCK IoStatus, 
							 PCOMPRESSED_DATA_INFO CompressedDataInfo, ULONG CompressedDataInfoLength, PDEVICE_OBJECT DeviceObject);
BOOLEAN FastIoWriteCompressed(PFILE_OBJECT FileObject, PLARGE_INTEGER FileOffset, ULONG Length,
						  ULONG LockKey, PVOID Buffer, PMDL *MdlChain, PIO_STATUS_BLOCK IoStatus,
						  PCOMPRESSED_DATA_INFO CompressedDataInfo, ULONG CompressedDataInfoLength, PDEVICE_OBJECT DeviceObject);
BOOLEAN FastIoMdlReadCompleteCompressed(PFILE_OBJECT FileObject, PMDL MdlChain, PDEVICE_OBJECT DeviceObject);
BOOLEAN FastIoMdlWriteCompleteCompressed(PFILE_OBJECT FileObject, PLARGE_INTEGER FileOffset, PMDL MdlChain, PDEVICE_OBJECT DeviceObject);
BOOLEAN FastIoQueryOpen(PIRP Irp, PFILE_NETWORK_OPEN_INFORMATION NetworkInformation, PDEVICE_OBJECT DeviceObject);

VOID DriverFsNotification(PDEVICE_OBJECT DeviceObject, BOOLEAN Active);
PLOG_RECORD StartLogIrp(PIRP Irp);
VOID EndLogIrp(PIRP Irp, PLOG_RECORD LogRecord);
VOID LogFastIo(FASTIO_TYPE FastIoType, PDEVICE_OBJECT DeviceObject, PFILE_OBJECT FileObject, PLARGE_INTEGER FileOffset, ULONG Lenght, BOOLEAN Wait, NTSTATUS Status);
VOID GetProcessEXEName(IN  ULONG pid, OUT PWCHAR ExeName);
NTSTATUS VdoToVolumeLabel(PDEVICE_OBJECT Vdo, PWCHAR Label);
VOID GetVolumeNameWorker(PDEVICE_OBJECT DeviceObject, PVOID Context);
VOID AttachToMountedVolumeWorker(PDEVICE_OBJECT DeviceObject, PVOID Context);
NTSTATUS GetFullFileName(PFILE_OBJECT FileObject, PDEVICE_OBJECT LowDevice, PWCHAR FilePath);
PVOID PsGetProcessSectionBaseAddress(PEPROCESS pEprocess);

#define IOCTL_START_FILTERING CTL_CODE(FILE_DEVICE_UNKNOWN, 0x800, METHOD_BUFFERED, FILE_ANY_ACCESS)
#define IOCTL_STOP_FILTERING CTL_CODE(FILE_DEVICE_UNKNOWN, 0x801, METHOD_BUFFERED, FILE_ANY_ACCESS)
#define IOCTL_GET_RECORD CTL_CODE(FILE_DEVICE_UNKNOWN, 0x803, METHOD_NEITHER, FILE_ANY_ACCESS)
#define IOCTL_ADD_FILTER_VALUE CTL_CODE(FILE_DEVICE_UNKNOWN, 0x804, METHOD_BUFFERED, FILE_ANY_ACCESS)
#define IOCTL_GET_COUNT_RECORD CTL_CODE(FILE_DEVICE_UNKNOWN, 0x805, METHOD_BUFFERED, FILE_ANY_ACCESS)
#define IOCTL_DISABLE_FILTER CTL_CODE(FILE_DEVICE_UNKNOWN, 0x806, METHOD_BUFFERED, FILE_ANY_ACCESS)